#include "Pch.h"
#include <string.h>
#include <math.h>

#include "PDBManager.h"
#include "MainWindow.h"
#include "Config.h"
#include "Atom.h"

extern "C" {
	#include "Matt/OctTree.h"
}

void Transform::CalcFullCurrentRotation(Matrix *out, PDBManager *pdbManager) {
	Matrix temp;
	Vector offsetVect;
	transformVect(&offsetVect, &worldView, &pdbManager->selCenter);
	Matrix offset = identity;
	Matrix unoffset = identity;
	for (int i=0; i<3; i++) {
		offset.M[i][3] = -offsetVect.v[i];
		unoffset.M[i][3] = offsetVect.v[i];
	}
	mulMat(&temp, &currentRot, &offset);
	mulMat(out, &unoffset, &temp);
}

void Transform::KeepCurrentRotation(PDBManager *pdbManager) {
	Matrix temp = worldView;
	Matrix rotation;
	CalcFullCurrentRotation(&rotation, pdbManager);
	mulMat(&worldView, &rotation, &temp);
	for (int i=0; i<3; i++) {
		translation.v[i] += worldView.M[i][3];
		worldView.M[i][3] = 0;
	}
	currentRot = identity;
}


Color PDBManager::GetResidueColor(int c, int index, int ignoreSelection) {
	if (index < 0) {
		return config->secondaryColors[0];
	}

	ManagedChain *chain = chains[c];
	unsigned int residueFlags = chain->residueInfo[index].flags;
	// ???
	// if (!(residueFlags & RESIDUE_RENDER)) return 0;
	if (!ignoreSelection) {
		if (residueFlags & (RESIDUE_SELECTED | RESIDUE_LAST_SELECTED)) {
			if (residueFlags & RESIDUE_LAST_SELECTED) {
				return config->colorSelRecent;
			}
			else {
				return config->colorSel;
			}
		}
	}
	if (!(residueFlags & RESIDUE_FORCED_COLOR)) {
		return chain->residueInfo[index].color;
	}
	else if (residueFlags & RESIDUE_COLOR_STRUCTURE) {
		int index2 = 0;
		if (residueFlags & RESIDUE_IS_ALPHA)
			index2 = 1;
		if (residueFlags & RESIDUE_IS_BETA)
			index2 = 2;
		if (residueFlags & RESIDUE_IS_DNA)
			index2 = 3;
		return config->secondaryColors[index2];
	}
	else if (residueFlags & RESIDUE_COLOR_ALIGNMENT) {
		const static Color colors[] = {
			{255,  0,  0},
			{  0,176,  0},
			{  0,  0,255},
			{176,176,  0},
			{  0,176,176},
			{176,  0,176},
		};
		const static Color unAligned = {255, 255, 255};
		if (!(residueFlags & RESIDUE_ALIGNED)) return unAligned;
		return colors[c % (sizeof(colors)/sizeof(colors[0]))];
	}
	// Should never be reached.
	return config->secondaryColors[0];
}

PDBManager::PDBManager(MainWindow *window) :
models(0),
numModels(0),
chains(0),
numChains(0),
scale(1),
undo(0),
redo(0),
windows(0),
numWindows(0),
align(0)
{
	AddWindow(window);
}

void PDBManager::AddWindow(MainWindow *window) {
	for (int i=0; i<numWindows; i++) {
		// By design, will happen once just after creation.
		// Will also happy when reusing it.
		if (windows[i] == window)
			return;
	}
	windows = (MainWindow**) realloc(windows, sizeof(MainWindow*) * (numWindows+1));
	windows[numWindows++] = window;
}

int PDBManager::RemoveWindow(MainWindow *window) {
	for (int i=0; i<numWindows; i++) {
		if (windows[i] == window) {
			windows[i] = windows[numWindows-1];
			windows = (MainWindow**) realloc(windows, sizeof(MainWindow*) * (numWindows-1));
			numWindows --;
			if (!numWindows) {
				delete this;
				return 0;
			}
			return numWindows;
		}
	}
	// ????
	return -1;
}

PDBManager::~PDBManager() {
	ClearPDB();
	// Should be null, anways.
	if (windows) free(windows);
}

void PDBManager::ClearPDB() {
	if (align) {
		CleanupSequenceIndexAlignment(align);
		align = 0;
	}
	ClearUndo();
	ClearRedo();

	// Have to get rid of chain stuff first -
	// CleanupModel call is actually responsible for freeing
	// everything else.
	for (int i=0; i<numChains; i++) {
		free(chains[i]->atomInfo);
		free(chains[i]->connectionInfo);
		free(chains[i]->residueInfo);
	}
	free(chains);
	chains = 0;
	numChains = 0;

	for (int i=0; i<numModels; i++) {
		free(models[i]->hBonds);
		CleanupModel(models[i]);
	}
	free(models);
	models = 0;
	numModels = 0;

	scale = 1;
}

int PDBManager::LoadAlignment(MultipleAlignment *ma) {
	ClearPDB();
	align = GetSequenceAlignmentIndices(ma);
	ManagedModel *model = 0;

	models = (ManagedModel**) realloc(models, sizeof(ManagedModel*) * (numModels+1));
	model = models[numModels] = (ManagedModel*) malloc(sizeof(ManagedModel));

	model->number = 0;
	model->numChains = ma->numChains;
	model->pdbName = strdup("Alignment");
	model->chains = (PDBChain**)malloc(sizeof(PDBChain*) * ma->numChains);
	for (int i=0; i<ma->numChains; i++) {
		model->chains[i] = ma->chains[i]->pdb;
	}

	model->hBonds = 0;
	model->numHbonds = 0;
	int added = 0;

	for (int i=0; i<model->numChains; i++) {
		added += TryAddChain(model, i);
	}
	if (added) {
		numModels++;
	}
	else {
		free(model);
		model = 0;
	}
	for (int i=0; i<ma->numChains; i++) {
		free(ma->chains[i]->res);
		free(ma->chains[i]);
	}
	CleanupAlignment(ma);

	CalcMatrices();
	CalcConnections();
	CalcSelectionCenter();

	if (align) {
		for (int i=0; i<align->length; i++) {
			int aligned = 1;
			for (int c=0; c<numChains; c++) {
				if (align->indices[c][i] < 0) {
					aligned = 0;
					break;
				}
			}
			if (!aligned) continue;
			for (int c=0; c<numChains; c++) {
				chains[c]->residueInfo[align->indices[c][i]].flags |= RESIDUE_ALIGNED;
			}
		}
	}

	return 1;
}

int PDBManager::TryAddChain(ManagedModel *m, int i) {
	chains = (ManagedChain**) realloc(chains, sizeof(ManagedChain*) * (numChains+1));
	int initialized = 0;
	if (!m->chains[i]->numTotalAtoms) return 0;
	CalcChainSecondary(m->chains[i]);
	initialized = 1;
	ManagedChain * chain = chains[numChains] = (ManagedChain*)calloc(1, sizeof(ManagedChain));
	*(PDBChain*)chain = *m->chains[i];
	free(m->chains[i]);
	m->chains[i] = chain;
	chain->model = m;
	chain->numChainResidues = chain->length;

	int firstAtomIndex = 0;
	if (numChains)
		firstAtomIndex = chains[numChains-1]->lastAtomIndex;

	chain->firstAtomIndex = firstAtomIndex;
	chain->lastAtomIndex = firstAtomIndex += chain->numTotalAtoms;

	free(chain->residueInfo);
	chain->residueInfo = (ExtendedResidueInfo*) calloc(chain->numTotalResidues, sizeof(ExtendedResidueInfo));

	for (int r=0; r<chain->numChainResidues; r++) {
		PDBResidue *res = &chain->residues[r];
		const char *dnaRes[] = {
			"  A", "  T", "  G", "  C", "  U", "  I",
			 " +A", " +T", " +G", " +C", " +U", " +I",
			 " DA", " DT", " DG", " DC", " DU", " DI"
		};
		int w;
		for (w = sizeof(dnaRes)/sizeof(dnaRes[0])-1; w>=0; w--) {
			if (!strcmp(res->resName, dnaRes[w])) break;
		}
		if (w >= 0) {
			chain->residueInfo[r].flags |= RESIDUE_IS_DNA;
		}
		else {
			chain->residueInfo[r].flags |= RESIDUE_IS_RESIDUE;
		}
	}

	for (int q=0; q<chain->numAlphaHelices; q++) {
		AlphaHelix *helix = chain->alphaHelices+q;
		for (int j=0; j<helix->length; j++) {
			chain->residueInfo[j + helix->start].flags |= RESIDUE_IS_ALPHA;
		}
	}
	for (int q=0; q<chain->numBetaPairs; q++) {
		BetaPair *pair = chain->betaPairs+q;
		for (int j=0; j<pair->length; j++) {
			chain->residueInfo[pair->start1+j].flags |= RESIDUE_IS_BETA;
			chain->residueInfo[pair->start2+j*pair->direction].flags |= RESIDUE_IS_BETA;
		}
	}

	numChains++;
	return 1;
}

int PDBManager::LoadPDB(PDBData *pdb) {
	ManagedModel *model = 0;
	ClearPDB();
	int oldNumChains = numChains;
	if (pdb) {
		if (pdb->numModels && pdb->numModels && pdb->models[0]->numChains) {
			models = (ManagedModel**) realloc(models, sizeof(ManagedModel*) * (numModels+1));
			model = models[numModels] = (ManagedModel*)malloc(sizeof(ManagedModel));
			*(PDBModel*)model = *pdb->models[0];
			free(pdb->models[0]);
			pdb->models[0] = model;
			model->hBonds = 0;
			model->numHbonds = 0;

			int added = 0;
			for (int i=0; i<model->numChains; i++) {
				added += TryAddChain(model, i);
			}
			if (added) {
				numModels++;
				pdb->numModels--;
				pdb->models[0] = pdb->models[pdb->numModels];
			}
			else {
				model = 0;
			}
		}
		CleanupPDB(pdb);
	}
	CalcMatrices();
	CalcConnections();
	CalcSelectionCenter();
	if (model) return 0;
	return 1;
}


void PDBManager::CalcSelectionCenter() {
	unsigned int selected = 0;
	unsigned int nonHidden = 1;
	{
		unsigned int min, max;
		GetMinMaxFlags(0, &min, &max);
		nonHidden = max & RESIDUE_RENDER;
		if (nonHidden) {
			GetMinMaxFlags(0, &min, &max, 1);
			selected = max & RESIDUE_SELECTED;
		}
	}
	Vector min = {0,0,0};
	Vector max = {0,0,0};
	int initialized = 0;
	for (int m=0; m<numModels; m++) {
		ManagedModel *model = models[m];
		for (int c=0; c<numChains; c++) {
			ManagedChain *chain = chains[c];
			if (chain->model != model || !chain->numTotalAtoms) continue;
			for (int j=0; j<chain->numTotalAtoms; j++) {
				if (selected && !(chain->residueInfo[chain->atoms[j].resIndex].flags & selected)) continue;
				if (nonHidden && !(chain->residueInfo[chain->atoms[j].resIndex].flags & RESIDUE_RENDER)) continue;
				Vector v = chain->atoms[j].pos;
				if (!initialized) {
					initialized = 1;
					min = max = v;
				}
				if (j >= chain->numChainResidues && !strcmp(chain->residues[chain->atoms[j].resIndex].resName, "HOH")) {
					continue;
				}
				for (int k=0; k<3; k++) {
					if (v.v[k] < min.v[k]) {
						min.v[k] = v.v[k];
					}
					if (v.v[k] > max.v[k]) {
						max.v[k] = v.v[k];
					}
				}
			}
		}
	}
	for (int i=0; i<3; i++) selCenter.v[i] = (min.v[i] + max.v[i])/2;
}

void PDBManager::CalcMatrices() {
	for (int m=0; m<numModels; m++) {
		ManagedModel *model = models[m];
		int initialized = 0;
		Vector Min = {0,0,0}, Max = {0,0,0}, center;
		for (int c=0; c<numChains; c++) {
			ManagedChain *chain = chains[c];
			if (chain->model != model || !chain->numTotalAtoms) continue;
			if (!initialized) {
				if (!initialized) {
					initialized = 1;
					Min = chain->atoms[0].pos;
					Max = chain->atoms[0].pos;
				}
			}
			for (int j=0; j<chain->numTotalAtoms; j++) {
				Vector *v = &chain->atoms[j].pos;
				for (int k=0; k<3; k++) {
					if (v->v[k] < Min.v[k]) {
						Min.v[k] = v->v[k];
					}
					if (v->v[k] > Max.v[k]) {
						Max.v[k] = v->v[k];
					}
				}
			}
		}
		for (int j=0; j<3; j++) {
			Min.v[j] -= 0.1;
			Max.v[j] += 0.1;
		}
		for (int k=0; k<3; k++) {
			center.v[k] = (Max.v[k] + Min.v[k])/2;
		}
		subVect(&model->min, &Min, &center);
		subVect(&model->max, &Max, &center);
		for (int c=0; c<numChains; c++) {
			ManagedChain *chain = chains[c];
			if (chain->model != model || !chain->numTotalAtoms) continue;
			for (int j=0; j<chain->numTotalAtoms; j++) {
				subVect(&chain->atoms[j].pos, &chain->atoms[j].pos, &center);
			}
		}
	}


	double maxDist = 0;
	for (int m=0; m<numModels; m++) {
		ManagedModel *model = models[m];
		int initialized = 0;
		Vector Min = {0,0,0}, Max = {0,0,0};
		for (int c=0; c<numChains; c++) {
			ManagedChain *chain = chains[c];
			if (chain->model != model) continue;
			for (int j=0; j<chain->numTotalAtoms; j++) {
				double dist = lengthVect(&chain->atoms[j].pos);
				if (dist > maxDist) {
					maxDist = dist;
				}
			}
		}
	}

	// 3 angstrom border.
	maxDist += 3;

	scale = (float)(1/maxDist*0.923879);
}

void CalcCylFrom(mglMatrix *m, Vector *coords, Vector *coords2) {
	Vector d;
	subVect(&d, coords2, coords);

	Vector h = {0, 1, 0};
	//subVect(&d, coords2, coords);
	Matrix m1, m2, m3;

	// Adjust cylinder's height.
	m3 = identity;
	m3.M[1][1] = lengthVect(&d);

	// Adjust cylinder's angle.
	Quaternion q;
	normalizeVect(&d);
	double angle = acos(dotVect(&h, &d));
	mulVect(&q.v, normalizeVect(crossVect(&q.v, &d, &h)), sin(angle/2));
	q.w = cos(angle/2);
	quatToMat(&m2, &q);

	// Move cylinder.

	m1 = identity;
	m1.M[0][3] = coords->x;
	m1.M[1][3] = coords->y;
	m1.M[2][3] = coords->z;

	Matrix temp, temp2;
	mulMat(&temp, &m2, &m3);
	mulMat(&temp2, &m1, &temp);
	mglConvertMatrix(m, &temp2);
}

void PDBManager::CalcConnections(ManagedModel *model) {
	int a;
	OctTreeNode tree = {0};
	MemoryManagementInfo info = {0,0};
	InitOctTreeNode(&tree, &model->min, &model->max);
	int count = 0;
	for (int c=0; c<numChains; c++) {
		ManagedChain *chain = chains[c];
		if (chain->model != model) continue;
		count += chain->numTotalAtoms;
		for (a=0; a<chain->numTotalAtoms; a++) {
			AddEntry(&tree, &chain->atoms[a].pos, &info, 2*MAX_ATOMIC_RADIUS);
		}
	}

	free(model->hBonds);
	model->hBonds = 0;
	model->numHbonds = 0;

	PDBAtom **hits = (PDBAtom**) malloc(sizeof(PDBAtom*)*(count+MAX_OCT_TREE_VERTICES));
	for (int c=0; c<numChains; c++) {
		ManagedChain *chain = chains[c];
		if (chain->model != model) continue;
		free(chain->atomInfo);
		chain->atomInfo = (ExtendedAtomInfo*) calloc(chain->numTotalAtoms+1, sizeof(ExtendedAtomInfo));
		for (a=0; a<chain->numTotalAtoms; a++) {
			PDBAtom *atom = &chain->atoms[a];
			chain->atomInfo[a].atomicIntSymbol = AtomicSymbolToInt(atom->element);
		}

		model->hBonds = (HBondInfo*) realloc(model->hBonds, (model->numHbonds + chain->numHbonds) * sizeof(HBondInfo));
		for (int h=0; h<chain->numHbonds; h++) {
			HBondInfo *hb = &model->hBonds[model->numHbonds++];
			hb->chain1 = chain;
			hb->res1 = chain->hbonds[h].res1;
			hb->atom1 = chain->hbonds[h].atom1;
			hb->chain2 = chain;
			hb->res2 = chain->hbonds[h].res2;
			hb->atom2 = chain->hbonds[h].atom2;
		}
	}

	for (int c=0; c<numChains; c++) {
		ManagedChain *chain = chains[c];
		if (chain->model != model) continue;

		free(chain->connectionInfo);
		int numConnections = 0;
		for (a=0; a<chain->numTotalAtoms; a++) {
			PDBAtom *atom = &chain->atoms[a];

			chain->atomInfo[a].firstConnectionIndex = numConnections;

			Vector *coords = &atom->pos;
			double radius = GetAtomicAverageRadius(chain->atomInfo[a].atomicIntSymbol);

			int numHits = GetEntries2((Vector**)hits, (void**)(hits+count+MAX_OCT_TREE_VERTICES-1), &tree, coords, MAX_ATOMIC_RADIUS+radius);

			// Don't want atoms from small molecules.
			if (a < chain->numAtoms && !stricmp(atom->name, " CA ")) {
				chain->atoms[a].flags |= ATOM_BACKBONE;
				if (a > 0 && chain->atoms[a-1].res == atom->res && !stricmp(chain->atoms[a-1].name, " N  ")) {
					chain->atoms[a-1].flags |= ATOM_BACKBONE;
				}
				int a2 = a+1;
				if (a2 < chain->numAtoms && chain->atoms[a2].res == atom->res && !stricmp(chain->atoms[a2].name, " C  ")) {
					chain->atoms[a2].flags |= ATOM_BACKBONE;
					a2++;
				}
				if (a2 < chain->numAtoms && chain->atoms[a2].res == atom->res && !stricmp(chain->atoms[a2].name, " O  ")) {
					chain->atoms[a2].flags |= ATOM_BACKBONE;
					a2++;
				}
			}

			for (int i=0; i<numHits; i++) {
				if (hits[i] == atom) continue;
				Vector d;
				double distSq = lengthSquaredVect(subVect(&d, &hits[i]->pos, coords));
				if (distSq >= 4*MAX_ATOMIC_RADIUS*MAX_ATOMIC_RADIUS)
					continue;
				double atomicDiameter = (radius + GetAtomicAverageRadius(hits[i]->element));
				if (atomicDiameter*atomicDiameter < distSq)
					continue;
				if (numConnections % 4096 == 0)
					chain->connectionInfo = (ConnectionInfo*) realloc(chain->connectionInfo, sizeof(ConnectionInfo) * (numConnections + 4096));
				chain->connectionInfo[numConnections].destAtom = hits[i];

				numConnections++;
			}
		}
		chain->atomInfo[a].firstConnectionIndex = numConnections;
	}
	free(hits);
	FreeAllMemory(&info);
}

void PDBManager::CalcConnections() {
	for (int m=0; m<numModels; m++) {
		CalcConnections(models[m]);
	}
}

void PDBManager::GetMinMaxFlags(unsigned int reqFlags, unsigned int *min, unsigned int *max, int forceVisible) {
	*min = ~0;
	*max = 0;
	for (int c=0; c<numChains; c++) {
		ManagedChain *chain = chains[c];
		for (int i=0; i<chain->numTotalResidues; i++) {
			if (forceVisible && !(chain->residueInfo[i].flags & RESIDUE_RENDER))
				continue;
			if (reqFlags == (reqFlags & chain->residueInfo[i].flags)) {
				*min &= chain->residueInfo[i].flags;
				*max |= chain->residueInfo[i].flags;
			}
		}
	}
}

void PDBManager::UpdateResidues(unsigned int reqFlags, unsigned int newFlags, unsigned int clearFlags, int redraw) {
	UpdateResidues2(reqFlags, reqFlags, newFlags, clearFlags, redraw);
}

int PDBManager::UpdateResidues2(unsigned int reqFlagMask, unsigned int reqFlags, unsigned int newFlags, unsigned int clearFlags, int redraw) {
	int changed = 0;
	for (int c=0; c<numChains; c++) {
		ManagedChain *chain = chains[c];
		int i;
		for (i=0; i<chain->numTotalResidues; i++) {
			if ((chain->residueInfo[i].flags & reqFlagMask) != reqFlags)
				continue;
			unsigned int flags = (chain->residueInfo[i].flags & ~clearFlags) | newFlags;
			changed |= (chain->residueInfo[i].flags != flags);
			chain->residueInfo[i].flags = flags;
		}
	}
	if (redraw) Redraw(redraw);
	return changed;
}

void PDBManager::UpdateSelectedResidues(unsigned int newFlags, unsigned int clearFlags, int redraw) {
	unsigned int min, max;
	GetMinMaxFlags(0, &min, &max);
	UpdateResidues(max & RESIDUE_SELECTED, newFlags, clearFlags, redraw);
}

void PDBManager::UpdateSelectedResiduesColor(Color color, int redraw) {
	unsigned int min, max;
	GetMinMaxFlags(0, &min, &max);
	int selected = 0;
	if (max & RESIDUE_SELECTED) {
		selected = 1;
	}
	for (int c=0; c<numChains; c++) {
		ManagedChain *chain = chains[c];
		for (int i=0; i<chain->numTotalResidues; i++) {
			if (selected && !(chain->residueInfo[i].flags & RESIDUE_SELECTED))
				continue;
			chain->residueInfo[i].flags &= ~RESIDUE_FORCED_COLOR;
			chain->residueInfo[i].color = color;
		}
	}
	if (redraw) Redraw(redraw);
}

void PDBManager::ToggleSelectedFlags(unsigned int newFlags, unsigned int clearFlagsIfSet, unsigned int clearFlagsIfUnset, int redraw) {
	unsigned int min, max;
	GetMinMaxFlags(0, &min, &max);
	if (max & RESIDUE_SELECTED) {
		GetMinMaxFlags(RESIDUE_SELECTED, &min, &max);
	}

	if ((newFlags & min) == newFlags) {
		UpdateResidues(max & RESIDUE_SELECTED, 0, newFlags | clearFlagsIfUnset, redraw);
	}
	else {
		UpdateResidues(max & RESIDUE_SELECTED, newFlags, clearFlagsIfSet, redraw);
	}
}

void PDBManager::ClearUndo() {
	ClearUndoBuffer(undo);
}

void PDBManager::ClearRedo() {
	ClearUndoBuffer(redo);
}

void PDBManager::Undo() {
	SingleUndo(undo, redo);
}

void PDBManager::Redo() {
	SingleUndo(redo, undo);
}

void PDBManager::ClearUndoBuffer(UndoBuffer *&buffer) {
	while (buffer) {
		UndoBuffer *temp = buffer;
		buffer = buffer->prev;
		free(temp);
	}
}

UndoBuffer *PDBManager::MakeUndoState() {
	int resCount = 0;
	for (int i=0; i<numChains; i++) {
		resCount += chains[i]->numTotalResidues;
	}
	UndoBuffer *undo2 = (UndoBuffer*) malloc(sizeof(UndoBuffer) + sizeof(ResUndoInfo) * (resCount-1));
	undo2->prev = 0;
	int index = 0;
	for (int i=0; i<numChains; i++) {
		for (int r=0; r<chains[i]->numTotalResidues; r++) {
			ExtendedResidueInfo *resInfo = &chains[i]->residueInfo[r];
			ResUndoInfo * unInfo = &undo2->info[index];
			unInfo->flags = resInfo->flags;
			unInfo->color = resInfo->color;
			index++;
		}
	}
	return undo2;
}

void PDBManager::SingleUndo(UndoBuffer *&undo, UndoBuffer *&redo) {
	if (!undo) return;
	UndoBuffer *undo2 = MakeUndoState();
	int done = 0;
	while (undo && !done) {
		int index = 0;
		for (int i=0; i<numChains; i++) {
			for (int r=0; r<chains[i]->numTotalResidues; r++) {
				ExtendedResidueInfo *resInfo = &chains[i]->residueInfo[r];
				ResUndoInfo * unInfo = &undo->info[index++];
				done |= (resInfo->flags != unInfo->flags) | (resInfo->color.uVal != unInfo->color.uVal);
				resInfo->flags = unInfo->flags;
				resInfo->color = unInfo->color;
			}
		}
		// repeat until undoing actually does something.
		UndoBuffer *temp = undo;
		undo = undo->prev;
		free(temp);
	}
	// Not terribly efficient in some cases, but works
	if (done) {
		undo2->prev = redo;
		redo = undo2;
		undo2 = 0;
	}
	else {
		// If at end of undo buffer and didn't undo anything,
		// don't need to add a new undo state.
		free(undo2);
	}
	Redraw();
}

void PDBManager::SaveUndoState() {
	int resCount = 0;
	int reallySave = 0;
	if (!undo) reallySave = 1;
	else {
		int index = 0;
		for (int i=0; i<numChains; i++) {
			for (int r=0; r<chains[i]->numTotalResidues; r++) {
				ExtendedResidueInfo *resInfo = &chains[i]->residueInfo[r];
				ResUndoInfo * unInfo = &undo->info[index];
				reallySave |= (resInfo->flags != unInfo->flags) | (resInfo->color.uVal != unInfo->color.uVal);
				index++;
			}
		}
	}
	if (!reallySave) {
		return;
	}
	ClearRedo();
	UndoBuffer *undo2 = MakeUndoState();
	undo2->prev = undo;
	undo = undo2;
}

// Note:  Ignores bonus residues.
int PDBManager::IsChainSelected(int chain) {
	ManagedChain *c = chains[chain];
	for (int res=0; res<c->numChainResidues; res++) {
		if (!(c->residueInfo[res].flags & RESIDUE_SELECTED)) return 0;
	}
	return 1;
}

void PDBManager::SelectChain(int chain) {
	ManagedChain *c = chains[chain];
	// TODO:  Maybe should be numChainResidues?
	for (int res=0; res<c->numTotalResidues; res++) {
		c->residueInfo[res].flags = ((c->residueInfo[res].flags & ~RESIDUE_LAST_SELECTED) | RESIDUE_SELECTED);
	}
}

void PDBManager::UnselectChain(int chain) {
	ManagedChain *c = chains[chain];
	for (int res=0; res<c->numTotalResidues; res++) {
		c->residueInfo[res].flags = (c->residueInfo[res].flags & ~(RESIDUE_LAST_SELECTED | RESIDUE_SELECTED));
	}
}

void PDBManager::ToggleChainSelection(int chain) {
	if (IsChainSelected(chain)) {
		UnselectChain(chain);
	}
	else {
		SelectChain(chain);
	}
}

int PDBManager::GetChainModelIndex(int chain) {
	for (int m=0; m<numModels; m++) {
		if (models[m] == chains[chain]->model) return m;
	}
	return 0;
}

void PDBManager::Redraw(unsigned int flags, int chain, int residue) {
	if (flags & REDRAW_UPDATE_SELECTION) {
		CalcSelectionCenter();
	}
	static int base = 0;
	base ++;
	if (base > 10000) base = 0;
	for (int i=0; i<numWindows; i++) {
		windows[(i+base)%numWindows]->Redraw(flags, chain, residue);
	}
}

void PDBManager::ChangeSelection(int chain, int residue) {
	for (int i=0; i<numWindows; i++) {
		windows[i]->ChangeSelection(chain, residue);
	}
}

void PDBManager::InvertSelection() {
	SaveUndoState();
	for (int c=0; c<numChains; c++) {
		for (int r=0; r<chains[c]->numTotalResidues; r++) {
			chains[c]->residueInfo[r].flags ^= RESIDUE_SELECTED;
		}
	}
	Redraw();
}
